import numpy as np

import pandas as pd
import tensorflow as tf

from tensor2tensor.layers.common_attention import local_attention_2d

# Check that attention is commutative with splitting/concatenating v along the depth dimension
# Try varios split points
for split_idx in range(1, 30):

    batch, heads, length, depth_k, depth_v, query_shape = 3, 4, 25, 16, 30, (4, 4)

    q = tf.random_normal([batch, heads, length, length, depth_k], dtype=tf.float64)
    k = tf.random_normal([batch, heads, length, length, depth_k], dtype=tf.float64)
    v = tf.random_normal([batch, heads, length, length, depth_v], dtype=tf.float64)
    
    # Apply attention with the first part of v
    output_part1 = local_attention_2d(
        q,
        k,
        v[:, :, :, :, :split_idx],
        query_shape=query_shape,
        memory_flange=(3, 3))
    
    # Apply attention with the second part of v
    output_part2 = local_attention_2d(
        q,
        k,
        v[:, :, :, :, split_idx:],
        query_shape=query_shape,
        memory_flange=(3, 3))
    
    # Put together results of two parts
    output_concat = tf.concat([output_part1, output_part2], axis=4)
    
    # Apply attention with the original v
    output_full = local_attention_2d(
        q,
        k,
        v,
        query_shape=query_shape,
        memory_flange=(3, 3))
    
    # Compute the difference - should be small
    with tf.Session() as sess:
        res_diff = (output_concat - output_full).eval()
        print(np.abs(res_diff).max())


2.220446049250313e-15
2.6645352591003757e-15
2.220446049250313e-15
1.7763568394002505e-15
2.220446049250313e-15
2.220446049250313e-15
2.220446049250313e-15
0.0
2.6645352591003757e-15
1.7763568394002505e-15
2.220446049250313e-15
2.220446049250313e-15
2.6645352591003757e-15
2.220446049250313e-15
2.220446049250313e-15
0.0
2.6645352591003757e-15
2.220446049250313e-15
2.4424906541753444e-15
2.6645352591003757e-15
3.1086244689504383e-15
2.220446049250313e-15
2.220446049250313e-15
0.0
1.3322676295501878e-15
0.0
2.220446049250313e-15
0.0
1.7763568394002505e-15